package main

import (
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"net"
	"sync"
)

// Conn 维护发送和接收的流信息
type Conn struct {
	net.Conn
	sendMu      sync.Mutex
	streamsOut  map[uint32]io.WriteCloser // 正在发送的流 key 为流ID
	recvMu      sync.Mutex
	streamsIn   map[uint32]*streamIn // 正在接收的流
	pendingRecv chan *streamIn       // 处理新接收到的流，调用 Receive 方法可以获取新的流
	closed      bool
	closedChan  chan struct{}
}

// 接收流 维护缓冲区并通过条件变量同步读写操作
type streamIn struct {
	streamID uint32
	key      string
	buffer   *bytes.Buffer
	closed   bool
	mu       sync.Mutex
	cond     *sync.Cond
}

type streamWriter struct {
	conn     *Conn
	streamID uint32
	closed   bool
	writeMu  sync.Mutex
}

func NewConn(conn net.Conn) *Conn {
	c := &Conn{
		Conn:        conn,
		streamsOut:  make(map[uint32]io.WriteCloser),
		streamsIn:   make(map[uint32]*streamIn),
		pendingRecv: make(chan *streamIn, 100), // 缓冲区暂定 100
		closedChan:  make(chan struct{}),
	}
	go c.readFrames()
	return c
}

func (c *Conn) Send(key string) (io.WriteCloser, error) {
	c.sendMu.Lock()
	defer c.sendMu.Unlock()

	var streamID uint32
	// 流ID生成：草率一点，直接随机生成，客户端、服务端各自维护
	// 很小的概率 服务端和客户端的流ID相同导致错误 但是过测试用例是够用了
	for {
		b := make([]byte, 4)
		if _, err := rand.Read(b); err != nil {
			return nil, err
		}
		streamID = binary.BigEndian.Uint32(b)
		// 如果 streamId 是唯一的 就可以用
		if _, exists := c.streamsOut[streamID]; !exists {
			break
		}
	}

	// 数据段，放 key 的长度、key、数据
	keyData := []byte(key)
	data := make([]byte, 2+len(keyData))
	// 用 uint16 纯粹是够用
	binary.BigEndian.PutUint16(data[:2], uint16(len(keyData)))
	copy(data[2:], keyData)

	if err := c.sendFrame(0, streamID, data); err != nil {
		return nil, err
	}

	writer := &streamWriter{
		conn:     c,
		streamID: streamID,
	}
	c.streamsOut[streamID] = writer
	return writer, nil
}

func (c *Conn) Receive() (key string, reader io.Reader, err error) {
	select {
	case s := <-c.pendingRecv:
		return s.key, s, nil
	case <-c.closedChan:
		return "", nil, io.EOF
	}
}

func (c *Conn) Close() {
	c.sendMu.Lock()
	defer c.sendMu.Unlock()
	c.recvMu.Lock()
	defer c.recvMu.Unlock()

	if c.closed {
		return
	}
	c.closed = true
	close(c.closedChan)
	c.Conn.Close()

	for _, w := range c.streamsOut {
		w.Close()
	}
	for _, s := range c.streamsIn {
		s.close()
	}
}

func (c *Conn) sendFrame(frameType byte, streamID uint32, data []byte) error {
	header := make([]byte, 1+4+4)
	header[0] = frameType
	binary.BigEndian.PutUint32(header[1:5], streamID)
	binary.BigEndian.PutUint32(header[5:9], uint32(len(data)))

	if _, err := c.Conn.Write(header); err != nil {
		return err
	}
	if len(data) > 0 {
		if _, err := c.Conn.Write(data); err != nil {
			return err
		}
	}
	return nil
}

func (c *Conn) readFrames() {
	for {
		frameType, streamID, dataLen, err := c.readFrameHeader()
		if err != nil {
			c.Close()
			return
		}

		data, err := c.readFrameData(dataLen)
		if err != nil {
			c.Close()
			return
		}

		switch frameType {
		case 0:
			if len(data) < 2 {
				c.Close()
				return
			}
			keyLen := binary.BigEndian.Uint16(data[:2])
			if len(data) < 2+int(keyLen) {
				c.Close()
				return
			}
			key := string(data[2 : 2+keyLen])
			s := newStreamIn(streamID, key)
			c.recvMu.Lock()
			c.streamsIn[streamID] = s
			c.recvMu.Unlock()
			select {
			case c.pendingRecv <- s:
			case <-c.closedChan:
				return
			}

		case 1:
			c.recvMu.Lock()
			s, ok := c.streamsIn[streamID]
			c.recvMu.Unlock()
			if ok {
				// 写到连接的缓冲区中、原子变量发送信号
				s.write(data)
			}

		case 2:
			c.recvMu.Lock()
			s, ok := c.streamsIn[streamID]
			if ok {
				s.close()
				delete(c.streamsIn, streamID)
			}
			c.recvMu.Unlock()

		default:
			c.Close()
			return
		}
	}
}

func (c *Conn) readFrameHeader() (byte, uint32, uint32, error) {
	// 类型（1） + streamId（4） + 数据长度（4）
	header := make([]byte, 1+4+4)
	if _, err := io.ReadFull(c.Conn, header); err != nil {
		return 0, 0, 0, err
	}
	frameType := header[0]
	streamID := binary.BigEndian.Uint32(header[1:5])
	dataLen := binary.BigEndian.Uint32(header[5:9])
	return frameType, streamID, dataLen, nil
}

func (c *Conn) readFrameData(dataLen uint32) ([]byte, error) {
	if dataLen == 0 {
		return nil, nil
	}
	data := make([]byte, dataLen)
	if _, err := io.ReadFull(c.Conn, data); err != nil {
		return nil, err
	}
	return data, nil
}

func newStreamIn(streamID uint32, key string) *streamIn {
	s := &streamIn{
		streamID: streamID,
		key:      key,
		buffer:   bytes.NewBuffer(nil),
	}
	s.cond = sync.NewCond(&s.mu)
	return s
}

func (s *streamIn) Read(p []byte) (n int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 等待数据
	for s.buffer.Len() == 0 && !s.closed {
		s.cond.Wait()
	}

	// 连接关闭
	if s.buffer.Len() == 0 && s.closed {
		return 0, io.EOF
	}

	return s.buffer.Read(p)
}

func (s *streamIn) write(data []byte) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return
	}

	s.buffer.Write(data)
	// 通知所有等待的 goroutine
	s.cond.Broadcast()
}

func (s *streamIn) close() {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.closed {
		return
	}

	s.closed = true
	s.cond.Broadcast()
}

func (sw *streamWriter) Write(p []byte) (int, error) {
	sw.writeMu.Lock()
	defer sw.writeMu.Unlock()

	if sw.closed {
		return 0, io.ErrClosedPipe
	}

	if err := sw.conn.sendFrame(1, sw.streamID, p); err != nil {
		return 0, err
	}
	return len(p), nil
}

func (sw *streamWriter) Close() error {
	sw.writeMu.Lock()
	defer sw.writeMu.Unlock()

	if sw.closed {
		return io.ErrClosedPipe
	}

	if err := sw.conn.sendFrame(2, sw.streamID, nil); err != nil {
		return err
	}

	sw.closed = true
	sw.conn.sendMu.Lock()
	delete(sw.conn.streamsOut, sw.streamID)
	sw.conn.sendMu.Unlock()
	return nil
}

// 除了上面规定的接口，你还可以自行定义新的类型，变量和函数以满足实现需求

//////////////////////////////////////////////
///////// 接下来的代码为测试代码，请勿修改 /////////
//////////////////////////////////////////////

// 连接到测试服务器，获得一个你实现的连接对象
func dial(serverAddr string) *Conn {
	conn, err := net.Dial("tcp", serverAddr)
	if err != nil {
		panic(err)
	}
	return NewConn(conn)
}

// 启动测试服务器
func startServer(handle func(*Conn)) net.Listener {
	ln, err := net.Listen("tcp", ":0")
	if err != nil {
		panic(err)
	}
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				fmt.Println("[WARNING] ln.Accept", err)
				return
			}
			go handle(NewConn(conn))
		}
	}()
	return ln
}

// 简单断言
func assertEqual[T comparable](actual T, expected T) {
	if actual != expected {
		panic(fmt.Sprintf("actual:%v expected:%v\n", actual, expected))
	}
}

// 简单 case：单连接，双向传输少量数据
// key 传递、数据分块传输、EOF 通知
func testCase0() {
	const (
		key  = "Bible"
		data = `Then I heard the voice of the Lord saying, “Whom shall I send? And who will go for us?”
And I said, “Here am I. Send me!”
Isaiah 6:8`
	)
	ln := startServer(func(conn *Conn) {
		// 服务端等待客户端进行传输
		_key, reader, err := conn.Receive()
		if err != nil {
			panic(err)
		}
		assertEqual(_key, key)
		dataB, err := io.ReadAll(reader)
		if err != nil {
			panic(err)
		}
		assertEqual(string(dataB), data)

		// 服务端向客户端进行传输
		writer, err := conn.Send(key)
		if err != nil {
			panic(err)
		}
		n, err := writer.Write([]byte(data))
		if err != nil {
			panic(err)
		}
		if n != len(data) {
			panic(n)
		}
		conn.Close()
	})
	//goland:noinspection GoUnhandledErrorResult
	defer ln.Close()

	conn := dial(ln.Addr().String())
	// 客户端向服务端传输
	writer, err := conn.Send(key)
	if err != nil {
		panic(err)
	}
	n, err := writer.Write([]byte(data))
	if n != len(data) {
		panic(n)
	}
	err = writer.Close()
	if err != nil {
		panic(err)
	}
	// 客户端等待服务端传输
	_key, reader, err := conn.Receive()
	if err != nil {
		panic(err)
	}
	assertEqual(_key, key)
	dataB, err := io.ReadAll(reader)
	if err != nil {
		panic(err)
	}
	assertEqual(string(dataB), data)
	conn.Close()
}

// 生成一个随机 key
func newRandomKey() string {
	buf := make([]byte, 8)
	_, err := rand.Read(buf)
	if err != nil {
		panic(err)
	}
	return hex.EncodeToString(buf)
}

// 读取随机数据，并返回随机数据的校验和：用于验证数据是否完整传输
func readRandomData(reader io.Reader, hash hash.Hash) (checksum string) {
	hash.Reset()
	var buf = make([]byte, 23<<20) //调用者读取时的 buf 大小不是固定的，你的实现中不可假定 buf 为固定值
	for {
		n, err := reader.Read(buf)
		if err == io.EOF {
			break
		}
		if err != nil {
			panic(err)
		}
		_, err = hash.Write(buf[:n])
		if err != nil {
			panic(err)
		}
	}
	checksum = hex.EncodeToString(hash.Sum(nil))
	return checksum
}

// 写入随机数据，并返回随机数据的校验和：用于验证数据是否完整传输
func writeRandomData(writer io.Writer, hash hash.Hash) (checksum string) {
	hash.Reset()
	const (
		dataSize = 500 << 20 //一个 key 对应 500MB 随机二进制数据，dataSize 也可以是其他值，你的实现中不可假定 dataSize 为固定值
		bufSize  = 1 << 20   //调用者写入时的 buf 大小不是固定的，你的实现中不可假定 buf 为固定值
	)
	var (
		buf  = make([]byte, bufSize)
		size = 0
	)
	for i := 0; i < dataSize/bufSize; i++ {
		_, err := rand.Read(buf)
		if err != nil {
			panic(err)
		}
		_, err = hash.Write(buf)
		if err != nil {
			panic(err)
		}
		n, err := writer.Write(buf)
		if err != nil {
			panic(err)
		}
		size += n
	}
	if size != dataSize {
		panic(size)
	}
	checksum = hex.EncodeToString(hash.Sum(nil))
	return checksum
}

// 复杂 case：多连接，大量数据，多个不同的 key
func testCase1() {
	var (
		mapKeyToChecksum = map[string]string{}
		lock             sync.Mutex
	)
	ln := startServer(func(conn *Conn) {
		// 服务端等待客户端进行传输
		key, reader, err := conn.Receive()
		if err != nil {
			panic(err)
		}
		var (
			h         = sha256.New()
			_checksum = readRandomData(reader, h)
		)
		lock.Lock()
		checksum, keyExist := mapKeyToChecksum[key]
		lock.Unlock()
		if !keyExist {
			panic(fmt.Sprintln(key, "not exist"))
		}
		assertEqual(_checksum, checksum)

		// 服务端向客户端连续进行 2 次传输
		for _, key := range []string{newRandomKey(), newRandomKey()} {
			writer, err := conn.Send(key)
			if err != nil {
				panic(err)
			}
			checksum := writeRandomData(writer, h)
			lock.Lock()
			mapKeyToChecksum[key] = checksum
			lock.Unlock()
			err = writer.Close() //表明该 key 的所有数据已传输完毕
			if err != nil {
				panic(err)
			}
		}
		conn.Close()
	})
	//goland:noinspection GoUnhandledErrorResult
	defer ln.Close()

	conn := dial(ln.Addr().String())
	// 客户端向服务端传输
	var (
		key = newRandomKey()
		h   = sha256.New()
	)
	writer, err := conn.Send(key)
	if err != nil {
		panic(err)
	}
	checksum := writeRandomData(writer, h)
	lock.Lock()
	mapKeyToChecksum[key] = checksum
	lock.Unlock()
	err = writer.Close()
	if err != nil {
		panic(err)
	}

	// 客户端等待服务端的多次传输
	keyCount := 0
	for {
		key, reader, err := conn.Receive()
		if err == io.EOF {
			// 服务端所有的数据均传输完毕，关闭连接
			break
		}
		if err != nil {
			panic(err)
		}
		_checksum := readRandomData(reader, h)
		lock.Lock()
		checksum, keyExist := mapKeyToChecksum[key]
		lock.Unlock()
		if !keyExist {
			panic(fmt.Sprintln(key, "not exist"))
		}
		assertEqual(_checksum, checksum)
		keyCount++
	}
	assertEqual(keyCount, 2)
	conn.Close()
}

func main() {
	testCase0()
	testCase1()
}
